import numpy as np
import pyximport
import torch
from torch.autograd import Variable
from torch.nn.modules.module import Module

from ..functions.gridgen import AffineGridGenFunction

pyximport.install(setup_args={"include_dirs": np.get_include()}, reload_support=True)


class _AffineGridGen(Module):
    def __init__(self, height, width, lr=1, aux_loss=False):
        super(_AffineGridGen, self).__init__()
        self.height, self.width = height, width
        self.aux_loss = aux_loss
        self.f = AffineGridGenFunction(self.height, self.width, lr=lr)
        self.lr = lr

    def forward(self, input):
        # if not self.aux_loss:
        return self.f(input)
        # else:
        #     identity = torch.from_numpy(np.array([[1,0,0], [0,1,0]], dtype=np.float32))
        #     batch_identity = torch.zeros([input.size(0), 2,3])
        #     for i in range(input.size(0)):
        #         batch_identity[i] = identity
        #     batch_identity = Variable(batch_identity)
        #     loss = torch.mul(input - batch_identity, input - batch_identity)
        #     loss = torch.sum(loss,1)
        #     loss = torch.sum(loss,2)

        #       return self.f(input), loss.view(-1,1)


class CylinderGridGen(Module):
    def __init__(self, height, width, lr=1, aux_loss=False):
        super(CylinderGridGen, self).__init__()
        self.height, self.width = height, width
        self.aux_loss = aux_loss
        self.f = CylinderGridGenFunction(self.height, self.width, lr=lr)
        self.lr = lr

    def forward(self, input):

        if not self.aux_loss:
            return self.f(input)
        else:
            return self.f(input), torch.mul(input, input).view(-1, 1)


class AffineGridGenV2(Module):
    def __init__(self, height, width, lr=1, aux_loss=False):
        super(AffineGridGenV2, self).__init__()
        self.height, self.width = height, width
        self.aux_loss = aux_loss
        self.lr = lr

        self.grid = np.zeros([self.height, self.width, 3], dtype=np.float32)
        self.grid[:, :, 0] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.height), 0),
                repeats=self.width,
                axis=0,
            ).T,
            0,
        )
        self.grid[:, :, 1] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.width), 0),
                repeats=self.height,
                axis=0,
            ),
            0,
        )
        self.grid[:, :, 2] = np.ones([self.height, width])
        self.grid = torch.from_numpy(self.grid.astype(np.float32))

    def forward(self, input1):
        self.batchgrid = torch.zeros(torch.Size([input1.size(0)]) + self.grid.size())

        for i in range(input1.size(0)):
            self.batchgrid[i] = self.grid
        self.batchgrid = Variable(self.batchgrid)

        if input1.is_cuda:
            self.batchgrid = self.batchgrid.cuda()

        output = torch.bmm(
            self.batchgrid.view(-1, self.height * self.width, 3),
            torch.transpose(input1, 1, 2),
        ).view(-1, self.height, self.width, 2)

        return output


class CylinderGridGenV2(Module):
    def __init__(self, height, width, lr=1):
        super(CylinderGridGenV2, self).__init__()
        self.height, self.width = height, width
        self.lr = lr
        self.grid = np.zeros([self.height, self.width, 3], dtype=np.float32)
        self.grid[:, :, 0] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.height), 0),
                repeats=self.width,
                axis=0,
            ).T,
            0,
        )
        self.grid[:, :, 1] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.width), 0),
                repeats=self.height,
                axis=0,
            ),
            0,
        )
        self.grid[:, :, 2] = np.ones([self.height, width])
        self.grid = torch.from_numpy(self.grid.astype(np.float32))

    def forward(self, input):
        self.batchgrid = torch.zeros(torch.Size([input.size(0)]) + self.grid.size())
        # print(self.batchgrid.size())
        for i in range(input.size(0)):
            self.batchgrid[i, :, :, :] = self.grid
        self.batchgrid = Variable(self.batchgrid)

        # print(self.batchgrid.size())

        input_u = input.view(-1, 1, 1, 1).repeat(1, self.height, self.width, 1)
        # print(input_u.requires_grad, self.batchgrid)

        output0 = self.batchgrid[:, :, :, 0:1]
        output1 = torch.atan(
            torch.tan(
                np.pi
                / 2.0
                * (
                    self.batchgrid[:, :, :, 1:2]
                    + self.batchgrid[:, :, :, 2:] * input_u[:, :, :, :]
                )
            )
        ) / (np.pi / 2)
        # print(output0.size(), output1.size())

        output = torch.cat([output0, output1], 3)
        return output


class DenseAffineGridGen(Module):
    def __init__(self, height, width, lr=1, aux_loss=False):
        super(DenseAffineGridGen, self).__init__()
        self.height, self.width = height, width
        self.aux_loss = aux_loss
        self.lr = lr

        self.grid = np.zeros([self.height, self.width, 3], dtype=np.float32)
        self.grid[:, :, 0] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.height), 0),
                repeats=self.width,
                axis=0,
            ).T,
            0,
        )
        self.grid[:, :, 1] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.width), 0),
                repeats=self.height,
                axis=0,
            ),
            0,
        )
        self.grid[:, :, 2] = np.ones([self.height, width])
        self.grid = torch.from_numpy(self.grid.astype(np.float32))

    def forward(self, input1):
        self.batchgrid = torch.zeros(torch.Size([input1.size(0)]) + self.grid.size())

        for i in range(input1.size(0)):
            self.batchgrid[i] = self.grid

        self.batchgrid = Variable(self.batchgrid)
        # print self.batchgrid,  input1[:,:,:,0:3]
        # print self.batchgrid,  input1[:,:,:,4:6]
        x = torch.mul(self.batchgrid, input1[:, :, :, 0:3])
        y = torch.mul(self.batchgrid, input1[:, :, :, 3:6])

        output = torch.cat([torch.sum(x, 3), torch.sum(y, 3)], 3)
        return output


class DenseAffine3DGridGen(Module):
    def __init__(self, height, width, lr=1, aux_loss=False):
        super(DenseAffine3DGridGen, self).__init__()
        self.height, self.width = height, width
        self.aux_loss = aux_loss
        self.lr = lr

        self.grid = np.zeros([self.height, self.width, 3], dtype=np.float32)
        self.grid[:, :, 0] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.height), 0),
                repeats=self.width,
                axis=0,
            ).T,
            0,
        )
        self.grid[:, :, 1] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.width), 0),
                repeats=self.height,
                axis=0,
            ),
            0,
        )
        self.grid[:, :, 2] = np.ones([self.height, width])
        self.grid = torch.from_numpy(self.grid.astype(np.float32))

        self.theta = self.grid[:, :, 0] * np.pi / 2 + np.pi / 2
        self.phi = self.grid[:, :, 1] * np.pi

        self.x = torch.sin(self.theta) * torch.cos(self.phi)
        self.y = torch.sin(self.theta) * torch.sin(self.phi)
        self.z = torch.cos(self.theta)

        self.grid3d = torch.from_numpy(
            np.zeros([self.height, self.width, 4], dtype=np.float32)
        )

        self.grid3d[:, :, 0] = self.x
        self.grid3d[:, :, 1] = self.y
        self.grid3d[:, :, 2] = self.z
        self.grid3d[:, :, 3] = self.grid[:, :, 2]

    def forward(self, input1):
        self.batchgrid3d = torch.zeros(
            torch.Size([input1.size(0)]) + self.grid3d.size()
        )

        for i in range(input1.size(0)):
            self.batchgrid3d[i] = self.grid3d

        self.batchgrid3d = Variable(self.batchgrid3d)
        # print(self.batchgrid3d)

        x = torch.sum(torch.mul(self.batchgrid3d, input1[:, :, :, 0:4]), 3)
        y = torch.sum(torch.mul(self.batchgrid3d, input1[:, :, :, 4:8]), 3)
        z = torch.sum(torch.mul(self.batchgrid3d, input1[:, :, :, 8:]), 3)
        # print(x)
        r = torch.sqrt(x ** 2 + y ** 2 + z ** 2) + 1e-5

        # print(r)
        theta = torch.acos(z / r) / (np.pi / 2) - 1
        # phi = torch.atan(y/x)
        phi = torch.atan(y / (x + 1e-5)) + np.pi * x.lt(0).type(torch.FloatTensor) * (
            y.ge(0).type(torch.FloatTensor) - y.lt(0).type(torch.FloatTensor)
        )
        phi = phi / np.pi

        output = torch.cat([theta, phi], 3)

        return output


class DenseAffine3DGridGen_rotate(Module):
    def __init__(self, height, width, lr=1, aux_loss=False):
        super(DenseAffine3DGridGen_rotate, self).__init__()
        self.height, self.width = height, width
        self.aux_loss = aux_loss
        self.lr = lr

        self.grid = np.zeros([self.height, self.width, 3], dtype=np.float32)
        self.grid[:, :, 0] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.height), 0),
                repeats=self.width,
                axis=0,
            ).T,
            0,
        )
        self.grid[:, :, 1] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.width), 0),
                repeats=self.height,
                axis=0,
            ),
            0,
        )
        self.grid[:, :, 2] = np.ones([self.height, width])
        self.grid = torch.from_numpy(self.grid.astype(np.float32))

        self.theta = self.grid[:, :, 0] * np.pi / 2 + np.pi / 2
        self.phi = self.grid[:, :, 1] * np.pi

        self.x = torch.sin(self.theta) * torch.cos(self.phi)
        self.y = torch.sin(self.theta) * torch.sin(self.phi)
        self.z = torch.cos(self.theta)

        self.grid3d = torch.from_numpy(
            np.zeros([self.height, self.width, 4], dtype=np.float32)
        )

        self.grid3d[:, :, 0] = self.x
        self.grid3d[:, :, 1] = self.y
        self.grid3d[:, :, 2] = self.z
        self.grid3d[:, :, 3] = self.grid[:, :, 2]

    def forward(self, input1, input2):
        self.batchgrid3d = torch.zeros(
            torch.Size([input1.size(0)]) + self.grid3d.size()
        )

        for i in range(input1.size(0)):
            self.batchgrid3d[i] = self.grid3d

        self.batchgrid3d = Variable(self.batchgrid3d)

        self.batchgrid = torch.zeros(torch.Size([input1.size(0)]) + self.grid.size())

        for i in range(input1.size(0)):
            self.batchgrid[i] = self.grid

        self.batchgrid = Variable(self.batchgrid)

        # print(self.batchgrid3d)

        x = torch.sum(torch.mul(self.batchgrid3d, input1[:, :, :, 0:4]), 3)
        y = torch.sum(torch.mul(self.batchgrid3d, input1[:, :, :, 4:8]), 3)
        z = torch.sum(torch.mul(self.batchgrid3d, input1[:, :, :, 8:]), 3)
        # print(x)
        r = torch.sqrt(x ** 2 + y ** 2 + z ** 2) + 1e-5

        # print(r)
        theta = torch.acos(z / r) / (np.pi / 2) - 1
        # phi = torch.atan(y/x)
        phi = torch.atan(y / (x + 1e-5)) + np.pi * x.lt(0).type(torch.FloatTensor) * (
            y.ge(0).type(torch.FloatTensor) - y.lt(0).type(torch.FloatTensor)
        )
        phi = phi / np.pi

        input_u = input2.view(-1, 1, 1, 1).repeat(1, self.height, self.width, 1)

        output = torch.cat([theta, phi], 3)

        output1 = torch.atan(
            torch.tan(
                np.pi
                / 2.0
                * (
                    output[:, :, :, 1:2]
                    + self.batchgrid[:, :, :, 2:] * input_u[:, :, :, :]
                )
            )
        ) / (np.pi / 2)
        output2 = torch.cat([output[:, :, :, 0:1], output1], 3)

        return output2


class Depth3DGridGen(Module):
    def __init__(self, height, width, lr=1, aux_loss=False):
        super(Depth3DGridGen, self).__init__()
        self.height, self.width = height, width
        self.aux_loss = aux_loss
        self.lr = lr

        self.grid = np.zeros([self.height, self.width, 3], dtype=np.float32)
        self.grid[:, :, 0] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.height), 0),
                repeats=self.width,
                axis=0,
            ).T,
            0,
        )
        self.grid[:, :, 1] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.width), 0),
                repeats=self.height,
                axis=0,
            ),
            0,
        )
        self.grid[:, :, 2] = np.ones([self.height, width])
        self.grid = torch.from_numpy(self.grid.astype(np.float32))

        self.theta = self.grid[:, :, 0] * np.pi / 2 + np.pi / 2
        self.phi = self.grid[:, :, 1] * np.pi

        self.x = torch.sin(self.theta) * torch.cos(self.phi)
        self.y = torch.sin(self.theta) * torch.sin(self.phi)
        self.z = torch.cos(self.theta)

        self.grid3d = torch.from_numpy(
            np.zeros([self.height, self.width, 4], dtype=np.float32)
        )

        self.grid3d[:, :, 0] = self.x
        self.grid3d[:, :, 1] = self.y
        self.grid3d[:, :, 2] = self.z
        self.grid3d[:, :, 3] = self.grid[:, :, 2]

    def forward(self, depth, trans0, trans1, rotate):
        self.batchgrid3d = torch.zeros(torch.Size([depth.size(0)]) + self.grid3d.size())

        for i in range(depth.size(0)):
            self.batchgrid3d[i] = self.grid3d

        self.batchgrid3d = Variable(self.batchgrid3d)

        self.batchgrid = torch.zeros(torch.Size([depth.size(0)]) + self.grid.size())

        for i in range(depth.size(0)):
            self.batchgrid[i] = self.grid

        self.batchgrid = Variable(self.batchgrid)

        x = self.batchgrid3d[:, :, :, 0:1] * depth + trans0.view(-1, 1, 1, 1).repeat(
            1, self.height, self.width, 1
        )

        y = self.batchgrid3d[:, :, :, 1:2] * depth + trans1.view(-1, 1, 1, 1).repeat(
            1, self.height, self.width, 1
        )
        z = self.batchgrid3d[:, :, :, 2:3] * depth
        # print(x.size(), y.size(), z.size())
        r = torch.sqrt(x ** 2 + y ** 2 + z ** 2) + 1e-5

        # print(r)
        theta = torch.acos(z / r) / (np.pi / 2) - 1
        # phi = torch.atan(y/x)
        phi = torch.atan(y / (x + 1e-5)) + np.pi * x.lt(0).type(torch.FloatTensor) * (
            y.ge(0).type(torch.FloatTensor) - y.lt(0).type(torch.FloatTensor)
        )
        phi = phi / np.pi

        # print(theta.size(), phi.size())

        input_u = rotate.view(-1, 1, 1, 1).repeat(1, self.height, self.width, 1)

        output = torch.cat([theta, phi], 3)
        # print(output.size())

        output1 = torch.atan(
            torch.tan(
                np.pi
                / 2.0
                * (
                    output[:, :, :, 1:2]
                    + self.batchgrid[:, :, :, 2:] * input_u[:, :, :, :]
                )
            )
        ) / (np.pi / 2)
        output2 = torch.cat([output[:, :, :, 0:1], output1], 3)

        return output2


class Depth3DGridGen_with_mask(Module):
    def __init__(self, height, width, lr=1, aux_loss=False, ray_tracing=False):
        super(Depth3DGridGen_with_mask, self).__init__()
        self.height, self.width = height, width
        self.aux_loss = aux_loss
        self.lr = lr
        self.ray_tracing = ray_tracing

        self.grid = np.zeros([self.height, self.width, 3], dtype=np.float32)
        self.grid[:, :, 0] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.height), 0),
                repeats=self.width,
                axis=0,
            ).T,
            0,
        )
        self.grid[:, :, 1] = np.expand_dims(
            np.repeat(
                np.expand_dims(np.arange(-1, 1, 2.0 / self.width), 0),
                repeats=self.height,
                axis=0,
            ),
            0,
        )
        self.grid[:, :, 2] = np.ones([self.height, width])
        self.grid = torch.from_numpy(self.grid.astype(np.float32))

        self.theta = self.grid[:, :, 0] * np.pi / 2 + np.pi / 2
        self.phi = self.grid[:, :, 1] * np.pi

        self.x = torch.sin(self.theta) * torch.cos(self.phi)
        self.y = torch.sin(self.theta) * torch.sin(self.phi)
        self.z = torch.cos(self.theta)

        self.grid3d = torch.from_numpy(
            np.zeros([self.height, self.width, 4], dtype=np.float32)
        )

        self.grid3d[:, :, 0] = self.x
        self.grid3d[:, :, 1] = self.y
        self.grid3d[:, :, 2] = self.z
        self.grid3d[:, :, 3] = self.grid[:, :, 2]

    def forward(self, depth, trans0, trans1, rotate):
        self.batchgrid3d = torch.zeros(torch.Size([depth.size(0)]) + self.grid3d.size())

        for i in range(depth.size(0)):
            self.batchgrid3d[i] = self.grid3d

        self.batchgrid3d = Variable(self.batchgrid3d)

        self.batchgrid = torch.zeros(torch.Size([depth.size(0)]) + self.grid.size())

        for i in range(depth.size(0)):
            self.batchgrid[i] = self.grid

        self.batchgrid = Variable(self.batchgrid)

        if depth.is_cuda:
            self.batchgrid = self.batchgrid.cuda()
            self.batchgrid3d = self.batchgrid3d.cuda()

        x_ = self.batchgrid3d[:, :, :, 0:1] * depth + trans0.view(-1, 1, 1, 1).repeat(
            1, self.height, self.width, 1
        )

        y_ = self.batchgrid3d[:, :, :, 1:2] * depth + trans1.view(-1, 1, 1, 1).repeat(
            1, self.height, self.width, 1
        )
        z = self.batchgrid3d[:, :, :, 2:3] * depth
        # print(x.size(), y.size(), z.size())

        rotate_z = (
            rotate.view(-1, 1, 1, 1).repeat(1, self.height, self.width, 1) * np.pi
        )

        x = x_ * torch.cos(rotate_z) - y_ * torch.sin(rotate_z)
        y = x_ * torch.sin(rotate_z) + y_ * torch.cos(rotate_z)

        r = torch.sqrt(x ** 2 + y ** 2 + z ** 2) + 1e-5

        # print(r)
        theta = torch.acos(z / r) / (np.pi / 2) - 1
        # phi = torch.atan(y/x)

        if depth.is_cuda:
            phi = torch.atan(y / (x + 1e-5)) + np.pi * x.lt(0).type(
                torch.cuda.FloatTensor
            ) * (
                y.ge(0).type(torch.cuda.FloatTensor)
                - y.lt(0).type(torch.cuda.FloatTensor)
            )
        else:
            phi = torch.atan(y / (x + 1e-5)) + np.pi * x.lt(0).type(
                torch.FloatTensor
            ) * (y.ge(0).type(torch.FloatTensor) - y.lt(0).type(torch.FloatTensor))

        phi = phi / np.pi

        output = torch.cat([theta, phi], 3)
        return output
